# BSD 3-Clause License
#
# Copyright (c) 2017 xxxx
# All rights reserved.
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# ============================================================================
import os
from decimal import Decimal
import utility

import torch
import torch.nn.utils as utils
from tqdm import tqdm

from apex import amp
import shutil

import numpy as np



class Trainer:
    def __init__(self, args, loader, my_model, my_loss, ckp):
        self.args = args
        self.scale = args.scale

        self.ckp = ckp
        self.loader_train = loader.loader_train
        self.loader_test = loader.loader_test

        self.model = my_model
        self.loss = my_loss
        self.optimizer = utility.make_optimizer(args, self.model)
        self.checkpoint_performance = []
        self.checkpoint_time = []
        if args.amp:
            if args.use_npu:
                self.model = self.model.to(args.device)
            self.model, self.optimizer = amp.initialize(
                self.model,
                self.optimizer,
                opt_level=args.opt_level,
                loss_scale=args.loss_scale,
            )
        if self.args.load != "":
            self.optimizer.load(ckp.dir, epoch=len(ckp.log))

        self.error_last = 1e8

    def train(self):
        self.loss.step()
        epoch = self.optimizer.get_last_epoch() + 1
        lr = self.optimizer.get_lr()

        self.ckp.write_log(
            "[Epoch {}]\tLearning rate: {:.2e}".format(epoch, Decimal(lr))
        )
        self.loss.start_log()
        self.model.train()

        timer_data, timer_model = utility.timer(), utility.timer()
        # TEMP
        self.loader_train.dataset.set_scale(0)
        print("epoch num ", epoch)


        for batch, (lr, hr, _,) in enumerate(self.loader_train):
            epoch_timer = utility.AverageMeter()
            lr, hr = self.prepare(lr, hr)
            timer_data.hold()
            timer_model.tic()

            self.optimizer.zero_grad()
            sr = self.model(lr, 0)
            loss = self.loss(sr, hr)
            if self.args.amp:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()
            if self.args.gclip > 0:
                utils.clip_grad_value_(
                    self.model.parameters(), self.args.gclip)
            self.optimizer.step()


            if self.args.use_npu:
                torch.npu.synchronize()
            timer_model.hold()
            epoch_timer.update(timer_model.acc)


            if (batch + 1) % self.args.print_every == 0:
                self.ckp.write_log(
                    "[{}/{}]\t{}\t{:.1f}+{:.1f}s".format(
                        (batch + 1) * self.args.batch_size,
                        len(self.loader_train.dataset),
                        self.loss.display_loss(batch),
                        timer_model.release(),
                        timer_data.release(),
                    )
                )


            timer_data.tic()
        print("FPS:", 1000  / epoch_timer.avg)

        self.loss.end_log(len(self.loader_train))
        self.error_last = self.loss.log[-1, -1]
        self.optimizer.schedule()

    def test(self):
        torch.set_grad_enabled(False)

        epoch = self.optimizer.get_last_epoch()
        self.ckp.write_log("\nEvaluation:")
        self.ckp.add_log(torch.zeros(
            1, len(self.loader_test), len(self.scale)))
        self.model.eval()

        timer_test = utility.timer()
        if self.args.save_results:
            self.ckp.begin_background()
        for idx_data, d in enumerate(self.loader_test):
            for idx_scale, scale in enumerate(self.scale):
                d.dataset.set_scale(idx_scale)
                for lr, hr, filename in tqdm(d, ncols=80):
                    lr, hr = self.prepare(lr, hr)
                    sr = self.model(lr, idx_scale)
                    sr = utility.quantize(sr, self.args.rgb_range)
                    save_list = [sr]
                    self.ckp.log[-1, idx_data, idx_scale] += utility.calc_psnr(
                        sr, hr, scale, self.args.rgb_range, dataset=d
                    )
                    if self.args.save_gt:
                        save_list.extend([lr, hr])

                    if self.args.save_results:
                        self.ckp.save_results(d, filename[0], save_list, scale)


                self.ckp.log[-1, idx_data, idx_scale] /= len(d)


                best = self.ckp.log.max(0)
                self.ckp.write_log(
                    "[{} x{}]\tPSNR: {:.3f} (Best: {:.3f} @epoch {})".format(
                        d.dataset.name,
                        scale,
                        self.ckp.log[-1, idx_data, idx_scale],
                        best[0][idx_data, idx_scale],
                        best[1][idx_data, idx_scale] + 1,
                    )
                )

        self.ckp.write_log("Forward: {:.2f}s\n".format(timer_test.toc()))
        self.ckp.write_log("Saving...")

        if self.args.save_results:
            self.ckp.end_background()

        if not self.args.test_only:
            self.ckp.save(self, epoch, is_best=(best[1][0, 0] + 1 == epoch))

        self.ckp.write_log("Total: {:.2f}s\n".format(
            timer_test.toc()), refresh=True)

        if not self.args.test_only:
            self.checkpoint_performance.append(
                [self.ckp.log[-1, idx_data, idx_scale]])
            self.checkpoint_time.append(timer_test.acc)
            is_best = (best[1][0, 0] + 1 == epoch)
            self.save_checkpoint(
                self.args,
                epoch,
                self.checkpoint_performance,
                self.checkpoint_time,
                self.model,
                self.optimizer,
                is_best,
            )
        torch.set_grad_enabled(True)

    def save_checkpoint(
        self, opt, epoch, checkpoint_performance, checkpoint_time, model, optimizer, is_best
    ):
        checkpoint_path = os.path.join(self.ckp.dir, "ckpt_latest.pth")
        checkpoint = {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "epoch": epoch,
            "amp": amp.state_dict() if opt.amp else None,
            "best": np.array(checkpoint_performance)[:, 0].max(),
            "ck_performance": checkpoint_performance,
            "ck_time": checkpoint_time,
        }
        torch.save(checkpoint, checkpoint_path)
        if is_best:
            print(
                "=> This Epoch {} is the BEST: {}".format(
                    epoch, checkpoint_performance[-1][0]
                )
            )
            opt.best = checkpoint_performance[-1][0]
            shutil.copyfile(
                checkpoint_path,
                os.path.join(self.ckp.dir, "ckpt_best.pth"),
            )

    def prepare(self, *args):
        device = torch.device("cpu" if self.args.cpu else "cuda")
        if self.args.use_npu:
            device = self.args.device

        def _prepare(tensor):
            return tensor.to(device)

        return [_prepare(a) for a in args]

    def terminate(self):
        if self.args.test_only:
            self.test()
            return True
        else:
            epoch = self.optimizer.get_last_epoch() + 1
            return epoch >= self.args.epochs
